"""
交易经纪人模块，负责处理订单和仓位
"""
import pandas as pd
import numpy as np
from enum import Enum
from collections import deque
from loguru import logger
from ..utils import load_ticker_info


class OrderType(Enum):
    """订单类型"""
    MARKET = "MARKET"  # 市价单
    LIMIT = "LIMIT"    # 限价单


class OrderAction(Enum):
    """订单动作"""
    BUY = "BUY"        # 买入
    SELL = "SELL"      # 卖出


class OrderOffset(Enum):
    """订单开平"""
    OPEN = "OPEN"      # 开仓
    CLOSE = "CLOSE"    # 平仓

class Order:
    """订单类"""
    
    def __init__(self, instrument, order_type, action, offset, quantity, price=None):
        """
        初始化订单
        
        Parameters
        ----------
        instrument : str
            交易品种代码
        order_type : OrderType
            订单类型
        action : OrderAction
            订单动作（买入或卖出）
        quantity : float
            数量
        price : float, optional
            价格，对于限价单是必须的
        """
        self.instrument = instrument
        self.order_type = order_type
        self.action = action
        self.offset = offset
        self.quantity = quantity
        self.price = price
        self.created_time = None
        self.filled_time = None
        self.filled_price = None
        self.order_id = None
        self.commission = None
    
    def __str__(self):
        return (f"Order(id={self.order_id}, {self.instrument}, {self.action.value}, "
                f"{self.order_type.value}, {self.offset.value}, quantity={self.quantity}, "
                f"price={self.price})")


class Broker:
    """
    经纪人类，负责处理订单和仓位
    """
    
    def __init__(self, initial_cash=100000.0, commission=0.0, timeline=None, ticker_info_file=None):
        """
        初始化经纪人
        
        Parameters
        ----------
        initial_cash : float, optional
            初始现金，默认为100000.0
        commission : float, optional
            佣金比例，默认为0.0
        timeline : pandas.Series, optional
            时间序列，包含所有可能的交易时间点
        ticker_info_file : str, optional
            交易品种信息文件路径，如果提供则自动加载品种信息
        """
        self._initial_cash = initial_cash
        self._cash = initial_cash
        self._commission = commission
        self._order_id_counter = 0
        
        # 从文件加载交易品种信息
        self._ticker_info = None
        if ticker_info_file:
            self._ticker_info = load_ticker_info(ticker_info_file)
            logger.info(f"已加载{len(self._ticker_info)}个交易品种的信息")
        
        # 主控时间轴
        self._timeline = None
        self._current_time_idx = 0
        self.set_timeline(timeline)
        
        # 滑点设置
        self._slippage = {}  # instrument -> 滑点值的映射
        # 如果有交易品种信息，自动设置滑点
        if self._ticker_info is not None:
            for inst, info in self._ticker_info.iterrows():
                self.set_slippage(inst, info['slippage'])
        
        # 使用DataFrame存储持仓信息
        self._positions = pd.DataFrame(columns=[
            'instrument',    # 交易品种
            'quantity',     # 持仓数量
            'avg_price',    # 平均持仓价格
            'cost_basis',   # 持仓成本
            'realized_pnl', # 已实现盈亏
            'unrealized_pnl'# 未实现盈亏
        ]).set_index('instrument')
        
        # 使用deque存储待处理订单
        self._pending_orders = deque()
        
        # 使用列表分批存储已成交订单
        self._filled_orders_chunks = []
        self._current_chunk = pd.DataFrame(columns=[
            'order_id',     # 订单ID
            'instrument',   # 交易品种
            'order_type',   # 订单类型
            'action',       # 买卖方向
            'quantity',     # 数量
            'price',        # 委托价格
            'filled_price', # 成交价格
            'created_time', # 创建时间
            'filled_time',  # 成交时间
            'commission'    # 佣金
        ]).set_index('order_id')
        self._chunk_size = 1000  # 每1000条订单创建新的chunk
        
        # 权益曲线，使用Series存储
        self._equity = None
    
    def set_slippage(self, instrument, slippage):
        """
        设置品种的滑点
        
        Parameters
        ----------
        instrument : str
            交易品种代码
        slippage : float
            滑点值（绝对价格）
        """
        self._slippage[instrument] = slippage
    
    def get_slippage(self, instrument):
        """
        获取品种的滑点
        
        Parameters
        ----------
        instrument : str
            交易品种代码
            
        Returns
        -------
        float
            滑点值，如果未设置则返回0
        """
        return self._slippage.get(instrument, 0.0)
    
    def set_timeline(self, timeline):
        """
        设置主控时间轴
        
        Parameters
        ----------
        timeline : pandas.Series
            时间序列，包含所有可能的交易时间点
        """
        if timeline is None:
            return
        self._timeline = timeline
        self._current_time_idx = 0
        # 初始化权益曲线
        self._equity = pd.Series(index=timeline, dtype=float)
        self._equity.iloc[0] = self._initial_cash
    
    @property
    def current_time(self):
        """获取当前时间"""
        if self._timeline is None:
            return None
        return self._timeline.iloc[self._current_time_idx]
    
    @property
    def cash(self):
        """获取当前现金"""
        return self._cash
    
    @property
    def positions(self):
        """获取当前所有持仓"""
        return self._positions
    
    @property
    def trades(self):
        """获取所有已成交订单"""
        # 合并所有chunks和当前chunk
        if not self._filled_orders_chunks:
            return self._current_chunk
        return pd.concat([*self._filled_orders_chunks, self._current_chunk])
    
    @property
    def equity(self):
        """获取权益曲线"""
        return self._equity
    
    def order(self, instrument, quantity,  action, offset=OrderOffset.OPEN, order_type=OrderType.MARKET, price=None,):
        """
        下单
        
        Parameters
        ----------
        instrument : str
            交易品种代码
        quantity : float
            数量，正数表示买入，负数表示卖出
        order_type : OrderType, optional
            订单类型，默认为市价单
        price : float, optional
            价格，对于限价单是必须的
            
        Returns
        -------
        int
            订单ID，如果资金不足返回None
        """
        # 检查参数
        if order_type == OrderType.LIMIT and price is None:
            raise ValueError("限价单必须指定价格")
        
        # 检查资金是否足够
        if  offset == OrderOffset.OPEN:
            estimated_price = price if price is not None else self._get_last_price(instrument)
            if estimated_price is None:
                logger.error(f"无法获取{instrument}的价格信息，订单创建失败")
                return None
                
            estimated_cost = quantity * estimated_price * (1 + self._commission)
            # TODO 这个判断只是判断了一个，如果一个bar多个合约开仓，就可能有问题。需要在订单提交的时候预扣减，然后处理完了之后再把预扣减变为真实扣减
            if estimated_cost > self._cash:
                logger.warning(
                    f"资金不足，订单创建失败。所需资金: {estimated_cost:.2f}, "
                    f"可用资金: {self._cash:.2f}"
                )
                return None

        # 创建订单
        order = Order(instrument, order_type, action, offset, quantity, price)
        order.order_id = self._order_id_counter
        order.created_time = self.current_time
        
        # 添加到待处理订单队列
        self._pending_orders.append(order)
        logger.info(
            f"创建订单: {order.order_id}, {instrument}, {action.value}, "
            f"{order_type.value}, {offset.value}, 数量={quantity}, 价格={price}"
        )
        
        # 更新订单ID计数器
        self._order_id_counter += 1
        
        return order.order_id
    
    def _get_last_price(self, instrument):
        """获取品种最新价格"""
        if instrument not in self._positions.index:
            return None
        position = self._positions.loc[instrument]
        return position['avg_price'] if position['quantity'] != 0 else None
    
    def update(self):
        """
        更新经纪人状态
        
        Parameters
        ----------
        price_data : dict
            价格数据字典，键为品种代码，值为当前bar数据
        """
        # TODO 由设置的feed来更新，而不是从参数获取
        if self._timeline is None:
            raise RuntimeError("必须先设置时间轴")
        
        # 处理待成交订单
        while self._pending_orders:
            order = self._pending_orders[0]
            if order.instrument in price_data:
                bar = price_data[order.instrument]
                self._process_order(order, bar)
                self._pending_orders.popleft()
            else:
                break
        
        # 更新持仓的未实现盈亏
        total_equity = self._cash
        for instrument, position in self._positions.iterrows():
            if instrument in price_data and position['quantity'] != 0:
                bar = price_data[instrument]
                position['unrealized_pnl'] = (bar['close'] - position['avg_price']) * position['quantity']
                total_equity += position['cost_basis'] + position['unrealized_pnl']
        
        # 记录权益
        self._equity.iloc[self._current_time_idx] = total_equity
        
        # 更新时间指针
        self._current_time_idx += 1
    
    def _process_order(self, order, bar):
        """
        处理订单
        
        Parameters
        ----------
        order : Order
            订单对象
        bar : pandas.Series
            当前bar数据
        """
        # 确定成交价格
        if order.order_type == OrderType.MARKET:
            fill_price = bar['close']
        else:  # LIMIT
            fill_price = order.price
        
        # 应用滑点
        slippage = self.get_slippage(order.instrument)
        if order.action == OrderAction.BUY:
            fill_price += slippage
        else:  # SELL
            fill_price -= slippage
        
        # 计算佣金
        commission = order.quantity * fill_price * self._commission
        
        # 更新订单信息
        order.filled_price = fill_price
        order.filled_time = self.current_time
        order.commission = commission
        
        # 将订单添加到当前chunk
        order_data = {
            'instrument': order.instrument,
            'order_type': order.order_type.value,
            'action': order.action.value,
            'quantity': order.quantity,
            'price': order.price,
            'filled_price': order.filled_price,
            'created_time': order.created_time,
            'filled_time': order.filled_time,
            'commission': order.commission
        }
        self._current_chunk.loc[order.order_id] = order_data
        
        # 如果当前chunk已满，创建新的chunk
        if len(self._current_chunk) >= self._chunk_size:
            self._filled_orders_chunks.append(self._current_chunk)
            self._current_chunk = pd.DataFrame(columns=self._current_chunk.columns)
        
        # 更新持仓
        self._update_position(order, fill_price, commission)
    
    def _update_position(self, order, fill_price, commission):
        """
        更新持仓
        
        Parameters
        ----------
        order : Order
            订单对象
        fill_price : float
            成交价格
        commission : float
            佣金
        """
        instrument = order.instrument
        quantity = order.quantity
        if order.action == OrderAction.SELL:
            quantity = -quantity
        
        # 计算交易成本
        cost = abs(quantity) * fill_price + commission
        
        # 如果没有该品种的持仓，创建一个新的持仓记录
        if instrument not in self._positions.index:
            self._positions.loc[instrument] = {
                'quantity': 0,
                'avg_price': 0,
                'cost_basis': 0,
                'realized_pnl': 0,
                'unrealized_pnl': 0
            }
        
        position = self._positions.loc[instrument]
        old_quantity = position['quantity']
        old_avg_price = position['avg_price']
        
        # 更新持仓
        new_quantity = old_quantity + quantity
        
        if new_quantity == 0:
            # 完全平仓
            realized_pnl = (fill_price - old_avg_price) * old_quantity - commission
            # 记录平仓信息
            logger.info(
                f"平仓: {instrument}, 数量={abs(old_quantity)}, "
                f"成交价={fill_price:.2f}, 盈亏={realized_pnl:.2f}"
            )
            # 删除持仓记录
            self._positions.drop(instrument, inplace=True)
            self._cash += fill_price * abs(old_quantity) - commission
        elif (old_quantity >= 0 and quantity > 0) or (old_quantity <= 0 and quantity < 0):
            # 同向加仓
            new_avg_price = (old_quantity * old_avg_price + quantity * fill_price) / new_quantity
            self._positions.loc[instrument, 'quantity'] = new_quantity
            self._positions.loc[instrument, 'avg_price'] = new_avg_price
            self._positions.loc[instrument, 'cost_basis'] += cost
            self._cash -= cost
            logger.info(
                f"加仓: {instrument}, 数量={abs(quantity)}, "
                f"均价={new_avg_price:.2f}, 总持仓={new_quantity}"
            )
        else:
            # 反向开仓或部分平仓
            if abs(old_quantity) > abs(quantity):
                # 部分平仓
                realized_pnl = (fill_price - old_avg_price) * quantity - commission
                self._positions.loc[instrument, 'quantity'] = new_quantity
                self._positions.loc[instrument, 'cost_basis'] = old_avg_price * abs(new_quantity)
                self._positions.loc[instrument, 'realized_pnl'] += realized_pnl
                self._cash += fill_price * abs(quantity) - commission
                logger.info(
                    f"部分平仓: {instrument}, 数量={abs(quantity)}, "
                    f"成交价={fill_price:.2f}, 盈亏={realized_pnl:.2f}, "
                    f"剩余持仓={new_quantity}"
                )
            else:
                # 完全平仓后反向开仓
                realized_pnl = (fill_price - old_avg_price) * old_quantity - commission
                self._positions.loc[instrument, 'quantity'] = new_quantity
                self._positions.loc[instrument, 'avg_price'] = fill_price
                self._positions.loc[instrument, 'cost_basis'] = abs(new_quantity) * fill_price
                self._positions.loc[instrument, 'realized_pnl'] += realized_pnl
                self._cash += fill_price * abs(old_quantity) - cost
                logger.info(
                    f"平仓并反向开仓: {instrument}, 平仓数量={abs(old_quantity)}, "
                    f"开仓数量={abs(new_quantity)}, 成交价={fill_price:.2f}, "
                    f"平仓盈亏={realized_pnl:.2f}"
                )
